# -*- coding: utf-8 -*-
"""
Created on Sun Jan  7 22:48:48 2018

@author: Luther
"""

print("功能初始化中......")

import jieba
import os
import re
import pandas as pd

df = pd.DataFrame()
year = []
codes = []

chn = re.compile(r'[\u4e00-\u9fa5]')


def tf(path, keyword_list, lengthtype):
    result = []
    try:
        with open(path, "r") as f:
            text = f.read()
    except:
        with open(path, "r", encoding='utf-8') as f:
            text = f.read()
    if lengthtype == 1:
        total = len(text)
    elif lengthtype == 2:
        l = jieba.cut(text)
        wl = [x for x in l if chn.match(x) and x not in stop_list]
        total = len(wl)
    else:
        raise ValueError("only 0 or 1 is permitted for lengthtype!")
    for keyword in keyword_list:
        k_num = text.count(keyword)
        result.append(str(k_num))
    result.append(str(total))
    return result


workpath = os.getcwd()
stop_path = os.path.join(workpath, "stop.txt")
with open(stop_path, "r") as f:
    stop = f.read()
stop_list = stop.split("\n")
print("{:}个停用词读取成功!".format(len(stop_list)))

keyword_filename = input("请输入关键词列表文件名：")
keyword_path = os.path.join(workpath, keyword_filename)
with open(keyword_path, "r") as f:
    keywords = f.read()
keyword_list = keywords.split("\n")
print("{:}个关键词读取成功!".format(len(keyword_list)))

print("功能初始化完成！")

path = input("请输入待检索文件目录：")
file_list = [x for x in os.listdir(path) if os.path.splitext(x)[-1] == '.txt']
full_path = [os.path.join(path, x) for x in file_list]
print("{:}个待检索文件读取成功!".format(len(file_list)))

lengthtype = int(input("全文长度：1；分词长度：2："))
print("正在检索......")

results = []
count = 0
for each in file_list:
    code = each[5:11]
    codes.append(code)
    y = each[:4]
    year.append(y)
    file_path = os.path.join(path, each)
    try:
        result = "/".join(tf(file_path, keyword_list, lengthtype))
    except:
        result = "0"
    results.append(result)
    count = count + 1
    print("\r当前进度: {:.2f}%".format(count * 100 / len(file_list)), end="")

df["code"] = codes
df["year"] = year
df["result"] = results
df["keyword"] = "/".join(keyword_list)
df["path"] = full_path

df.to_csv("results.csv")
output_path = input("请输入结果输出位置：")
df.to_csv(output_path)

print("完成！")
